import json
from typing import List
from os.path import abspath, join
from backend.experiment.clue.afqmc.scheme6.implement import AFQMCImplement
from backend.experiment.tokenizer import tokenizer
from backend.experiment.model import vocab_name
from backend.experiment.framework.shots_prompt.implement import ShotMode
from backend.experiment.framework.task import PerplexityChoiceMakingTask

test_dev = True
insert_start_token = True
only_calculate_last = True
seed = 42
task_handler = AFQMCImplement.get_task_handler()


class AFQMCSample:
    def __init__(self, s1: str, s2: str, label: str, task):
        self.__s1 = s1
        self.__s2 = s2
        self.__label = label
        self.__task = task

    @property
    def s1(self) -> str:
        return self.__s1

    @property
    def s2(self) -> str:
        return self.__s2

    @property
    def label(self) -> str:
        return self.__label

    @property
    def task(self) -> PerplexityChoiceMakingTask:
        return self.__task


def initialize() -> List[AFQMCSample]:
    with open(abspath(
            join(__file__, '../../experiment/clue/afqmc/dev.json')),
            'r', encoding='utf8') as file:
        records = [json.loads(each) for each in file.readlines()]

    cache_filename = f'{vocab_name}_one_shot_{seed}'
    if insert_start_token:
        cache_filename += '_start_token'

    if only_calculate_last:
        cache_filename += '_only_last_shot'

    test_tasks = AFQMCImplement(
        tokenizer=tokenizer,
        test_dev=test_dev,
        seed=seed
    ).get_tasks(
        mode=ShotMode.FEW_SHOT,
        cache_filename=cache_filename + '.json',
        insert_start_token=insert_start_token,
        only_calculate_last=only_calculate_last
    )

    _samples = []
    i = 0
    for each in records:
        if test_tasks[i] is not None:
            _samples.append(
                AFQMCSample(
                    s1=each['sentence1'],
                    s2=each['sentence2'],
                    label=AFQMCImplement.label_map()[int(each['label'])],
                    task=test_tasks[i]
                ))
        i += 1
    return _samples


samples = initialize()
